package edu.washington.ling.syntric.util;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import edu.stanford.nlp.util.StringUtils;
import edu.washington.ling.syntric.struct.error.IrregularTreeException;
import edu.washington.ling.syntric.struct.error.UnreadableRepresentationException;
import edu.washington.ling.syntric.struct.tree.*;


public class TreeUtils {
	
	public static void derriveDSLinksFromFeatures(Tree tree) throws UnreadableRepresentationException {
		for (Token t: tree.getTokenList()) {
			if (t.hasFeature("drel")) {
				String relParts[] = t.getValue("drel").split(":");
				t.setDSRelation(relParts[0]);
				if (!tree.hasTokenWithName(relParts[1])) {
					String errStr = String.format("drel specifies unknown name '%s'.", relParts[1]);
					throw new UnreadableRepresentationException(errStr);
				}
				Token parent = tree.getTokenWithName(relParts[1]);
				t.setDSHead(parent);
				parent.addDSDependent(t);
			} 
			else if (t.hasFeature("dmrel")) {
				String relParts[] = t.getValue("dmrel").split(":");
				t.setDSRelation(relParts[0]);
				if (!tree.hasTokenWithName(relParts[1])) {
					String errStr = String.format("dmrel specifies unknown name '%s'", relParts[1]);
					throw new UnreadableRepresentationException(errStr);
				}
				Token parent = tree.getTokenWithName(relParts[1]);
				t.setDSHead(parent);
				parent.addDSDependent(t);
			}
			else {
				tree.addDSRoot(t);
			}
			if (t.hasName() == false)  {
				String errStr = String.format("name attribute missing for word '%s'", t.getWord());
				throw new UnreadableRepresentationException(errStr);
			}
		}
	}
	public static void writeDSLinksAsFeatures(Tree tree) {
		for (Token t: tree.getTokenList()) {
			if (t.hasDSRelation() && t.hasDSHead()) {
				String val = String.format("%s:%s",t.getDSRelation(), t.getDSHead().getName());
				if (t.isEmptyCategory()) {
					t.setFeature("dmrel", val);
				}
				else {
					t.setFeature("drel", val);
				}
			}
		}
	}
	public static void derrivePBLinksFromFeatures(Tree tree) throws UnreadableRepresentationException {
		for (Token t: tree.getTokenList()) {
			if (t.hasFeature("pbrel")) {
				String relParts[] = t.getValue("pbrel").split(":");
				for (int i = 0; i < relParts.length; i+=2) {
					if (!tree.hasTokenWithName(relParts[i+1])) {
						String errStr = String.format("pbrel specifies unknown name '%s'.", relParts[i+1]);
						throw new UnreadableRepresentationException(errStr);
					}
					Token head = tree.getTokenWithName(relParts[i+1]);
					t.addPBHead(head, relParts[i]);
					head.addPBArgument(t);
				}
			} 
			if (t.hasFeature("pbmrel")) {
				String relParts[] = t.getValue("pbmrel").split(":");
				for (int i = 0; i < relParts.length; i+=2) {
					if (!tree.hasTokenWithName(relParts[i+1])) {
						String errStr = String.format("pbmrel specifies unknown name '%s'.", relParts[i+1]);
						throw new UnreadableRepresentationException(errStr);
					}
					Token head = tree.getTokenWithName(relParts[i+1]);
					t.addPBHead(head, relParts[i]);
					head.addPBArgument(t);
				}
			}
			if (t.hasName() == false)  {
				String errStr = String.format("name attribute missing for word '%s'", t.getWord());
				throw new UnreadableRepresentationException(errStr);
			}
		}
	}
	public static void writePBLinksAsFeatures(Tree tree) {
		for (Token t: tree.getTokenList()) {
			List<String> vals = new LinkedList<String>();
			for (Token head: t.getPBHeads()) {
				String argumentType = t.getPBArgumentType(head);
				vals.add(String.format("%s:%s",argumentType, head.getName()));
			}
			if (vals.size() > 0) {
				String val = StringUtils.join(vals, ":");
				if (t.isEmptyCategory()) {
					t.setFeature("pbmrel", val);
				}
				else {
					t.setFeature("pbrel", val);
				}
			}
		}
	}
	public static Pair<Integer, Integer> getSpan(Token t, Tree ds) {
		int leftBoundary = ds.getTokenPosition(t);
		int rightBoundary = ds.getTokenPosition(t);
		if (t.hasDSDependents()) {
			for (Token child: t.getDSDependents()) {
				Pair<Integer, Integer> childSpan = getSpan(child, ds);
				if (childSpan.getLeft() < leftBoundary) {
					leftBoundary = childSpan.getLeft();
				}
				if (childSpan.getRight() > rightBoundary) {
					rightBoundary = childSpan.getRight();
				}
			}
		}
		return new Pair<Integer, Integer>(leftBoundary, rightBoundary);
	}
	
	public static List<Token> getLeftDSChildren(Token t, Tree s) {
		List<Token> children = t.getDSDependents();
		List<Token> leftChildren = new LinkedList<Token>();
		int parentIdx = s.getTokenPosition(t);
		for (Token child: children) {
			if (s.getTokenPosition(child) < parentIdx) {
				leftChildren.add(child);
			} 
		}
		return leftChildren;
	}
	public static List<Token> getRightDSChildren(Token t, Tree s) {
		List<Token> children = t.getDSDependents();
		List<Token> rightChildren = new LinkedList<Token>();
		int parentIdx = s.getTokenPosition(t);
		for (Token child: children) {
			if (s.getTokenPosition(child) > parentIdx) {
				rightChildren.add(child);
			} 
		}
		return rightChildren;
	}
	public static List<Token> getLeftInnerDSSiblings(Token dep, Token head, Tree s) {
		List<Token> children = head.getDSDependents();
		List<Token> leftInnerSiblings = new LinkedList<Token>();
		int depPos = s.getTokenPosition(dep);
		int headPos = s.getTokenPosition(head);
		for (Token child: children) {
			int childPos = s.getTokenPosition(child);
			if (depPos < childPos && childPos < headPos) {
				leftInnerSiblings.add(child);
			} 
		}
		return leftInnerSiblings;
	}
	public static List<Token> getRightInnerDSSiblings(Token dep, Token head, Tree s) {
		List<Token> children = head.getDSDependents();
		List<Token> rightInnerSiblings = new LinkedList<Token>();
		int depPos = s.getTokenPosition(dep);
		int headPos = s.getTokenPosition(head);
		for (Token child: children) {
			int childPos = s.getTokenPosition(child);
			if (headPos < childPos && childPos < depPos) {
				rightInnerSiblings.add(child);
			} 
		}
		return rightInnerSiblings;
	}
	public static void sortTokens(List<Token> tokens, final Tree s) {
		sortTokens(tokens, s, false);
	}
	public static void sortTokens(List<Token> tokens, final Tree s, final boolean reverse) {
		Collections.sort(tokens, new Comparator<Token>() {
			public int compare(Token o1, Token o2) {
				if (reverse) {
					return s.getTokenPosition(o2) - s.getTokenPosition(o1);
				} else {
					return s.getTokenPosition(o1) - s.getTokenPosition(o2);
				}
			}
		});
	}
	private static void labelHeadWords(Token childToken, Token parentToken, String childHead, String parentHead) {
		TreeUtils.markLCA(childToken, parentToken);
		TreeUtils.markHeadsChildBranch(childToken, childHead);
		TreeUtils.markHeadsParentBranch(parentToken, parentHead);
		TreeUtils.unMarkLCA(childToken);
	}
	
	public static void labelHeadWordsForTree(Tree tree){
		for (Token t: tree.getTokenList()){
			if (t.hasDSHead()){
				Token childNode = t;
				Token parentNode = t.getDSHead();
				String childHead = childNode.getName();
				String parentHead = parentNode.getName();
				labelHeadWords(childNode, parentNode, childHead, parentHead);
			}
		}	
	}

	public static void touchBranch(Node n) {
		n.touch();
		if (n.hasPSParent()) {
			touchBranch(n.getPSParent());
		}
	}
	
	public static void untouchBranch(Node n) {
		n.unTouch();
		if (n.hasPSParent()) {
			untouchBranch(n.getPSParent());
		}
	}
	
	public static Node findFirstTouched(Node n) {
		if (n.isTouched())
			return n;
		else if (n.hasPSParent())
			return findFirstTouched(n.getPSParent());
		else
			return null;
	}
	
	public static Node markLCA(Node n1, Node n2){
		touchBranch(n1);
		Node lca = findFirstTouched(n2);
		if (lca != null)
			lca.markAsLCA();
		untouchBranch(n1);
		return lca;
	}
	
	public static void unMarkLCA(Node n){
		if (n.isLCA()) {
			n.unMarkAsLCA();
		} else if (n.hasPSParent()) {
			unMarkLCA(n.getPSParent());
		}
	}
	
	public static Node getChildOfLCA(Node n) {
		if (n.hasPSParent()) {
			if (n.getPSParent().isLCA()) {
				return n;
			} else {
				return getChildOfLCA(n.getPSParent());
			}
		} else {
			return null;
		}
	}
	
	public static void markHeadsChildBranch(Node n, String w) {
		n.addHeadWord(w);
		if (n.hasPSParent()) {
			if (!n.getPSParent().isLCA()){
				markHeadsChildBranch(n.getPSParent(), w);
			}
				
		}
	}
	
	public static void markHeadsParentBranch(Node n, String w) {
		n.addHeadWord(w);
		if (!n.isLCA() && n.hasPSParent()){
			markHeadsParentBranch(n.getPSParent(), w);
		}
	}
	
	public static void markHeadsMovedBranch(Node n1, Node n2, String w) {
		n1.addHeadWord(w);
		if (!n1.equals(n2) && n1.hasPSParent()) {
			markHeadsMovedBranch(n1.getPSParent(), n2, w);
		}
	}
	
	public static boolean isOnMovedBranch(Node n) {
		if (!n.isLCA() && n.hasFeature("index") && !(n instanceof Token)) {
			return true;
		}
		if (n.isLCA() || !n.hasPSParent()) {
			return false;
		}
		else {
			return isOnMovedBranch(n.getPSParent());
		}
	}
	public static Node getMovedBranchMP(Node n) {
		if (n.hasFeature("index")) {
			return n;
		}
		else {
			return getMovedBranchMP(n.getPSParent());
		}
	}
	public static Node findFirstBranching(Node n) {
		Node branchingNode = null;
		if (n.getPSChildren().size() > 1) {
			branchingNode = n;
		}
		else if (n.hasPSParent()) {
			branchingNode = findFirstBranching(n.getPSParent());
		}
		return branchingNode;
	}
	
	public static void checkDSRoot(Tree ds) throws IrregularTreeException {
		List<Token> dsRoots = ds.getDSRoots();
		ArrayList<String> subErrStrs = new ArrayList<String>();
		if (dsRoots.size() < 1) {
			String errStr = "Error: no DS root found.";
			subErrStrs.add(errStr);
		}
		else if (dsRoots.size() > 1) {
			for (Token dsRoot: dsRoots) {
				String errStr = String.format("Error: token \'%s\' is one of multiple DS roots.", dsRoot.getName());
				subErrStrs.add(errStr);
			}
		}
		if (subErrStrs.size() > 0) {
			String errStr = String.format("DS root is not unique");
			throw new IrregularTreeException(errStr, subErrStrs);
		}
	}
	public static List<Node> getAntecedents(Node node) {
		List<Node> antes = new LinkedList<Node>();
		if (!(node instanceof Token) && node.hasFeature("index")) {
			antes.add(node);
		}
		for (Node child: node.getPSChildren()) {
			antes.addAll(getAntecedents(child));
		}
		return antes;
	}
	public static Map<String, List<Node>> getAntecedentMap(Tree tree) {
		Map<String, List<Node>> anteMap = new HashMap<String, List<Node>>();
		List<Node> antes = getAntecedents(tree.getPSRoot());
		for (Node ante: antes) {
			String indexVal = ante.getValue("index");
			if (!anteMap.containsKey(indexVal)) {
				anteMap.put(indexVal, new LinkedList<Node>());	
			}
			anteMap.get(indexVal).add(ante);
		}
		return anteMap;
	}
	public static Map<String, List<Token>> getTraceMap(Tree tree) {
		Map<String, List<Token>> traceMap = new HashMap<String, List<Token>>();
		for (Token token: tree.getTokenList()) {
			if (token.hasFeature("index")) {
				Token trace = token;
				String indexVal = trace.getValue("index");
				if (!traceMap.containsKey(indexVal)) {
					traceMap.put(indexVal, new LinkedList<Token>());
				}
				traceMap.get(indexVal).add(trace);
			}
		}
		return traceMap;
	}
	public static void checkAntecedentsAndTraces(Tree tree) throws IrregularTreeException {
		Map<String, List<Node>> anteMap = getAntecedentMap(tree);
		Map<String, List<Token>> traceMap = getTraceMap(tree);
		ArrayList<String> subErrStrs = new ArrayList<String>();
		for (String index: traceMap.keySet()) {
			if (!anteMap.containsKey(index)) {
				String errStr = String.format("Error: missing antecedent with index \'%s\'", index);
				subErrStrs.add(errStr);
			} else {
				if (anteMap.get(index).size() > 1) {
					String errStr = String.format("Error: more than one antecedent with index \'%s\'", index);
					subErrStrs.add(errStr);
				}
			}
		}
		for (String index: anteMap.keySet()) {
			if (!traceMap.containsKey(index)) {
				String errStr = String.format("Error: missing trace for index \'%s\'", index);
				subErrStrs.add(errStr);
			} else {
				if (traceMap.get(index).size() > 1) {
					String errStr = String.format("Error: more than one trace with index \'%s\'", index);
					subErrStrs.add(errStr);
				}
			}
		}
		if (subErrStrs.size() > 0) {
			String errStr = String.format("Mismatched traces and antecedents found.");
			throw new IrregularTreeException(errStr, subErrStrs);
		}
	}
	
	public static void checkPreterminals(Tree ps) throws IrregularTreeException {
		ArrayList<String> subErrStrs = new ArrayList<String>();
		for (Token t: ps.getTokenList()) {
			if ((t.getPSParent().getPSChildren().size() > 1)) {
				String errStr = String.format("Error: preterminal for token \'%s\' has more than one child", t.getWord());
				subErrStrs.add(errStr);
			}
		}
		if (subErrStrs.size() > 0) {
			String errStr = String.format("Illegal preterminals found.");
			throw new IrregularTreeException(errStr, subErrStrs);
		}
	}
	
	private static boolean isNBToTerminal(Node n) {
		boolean isNBToTerminal = true;
		if(n.getPSChildren().size() > 1) {
			isNBToTerminal = false;
		}
		else if(n.getPSChildren().size() == 1) {
			isNBToTerminal = isNBToTerminal(n.getPSChildren().get(0));
		}
		else if(n.getPSChildren().size() == 0) {
			isNBToTerminal = true;
		}
		return isNBToTerminal;
	}
	private static List<Node> getNBInternalNodes(Node n) {
		List<Node> nbNodes = new LinkedList<Node>();
		if (n.getPSChildren().size() > 1) {
			for (Node child: n.getPSChildren()) {
				nbNodes.addAll(getNBInternalNodes(child));
			}
		}
		else if (n.getPSChildren().size() == 1) {
			if(!isNBToTerminal(n)) {
				nbNodes.add(n);
			}
			for (Node child: n.getPSChildren()) {
				nbNodes.addAll(getNBInternalNodes(child));
			}
		}
		return nbNodes;
	}
	
	public static void checkInternalNodeBranching(Tree ps) throws IrregularTreeException {
		List<Node> nbIntNodes = getNBInternalNodes(ps.getPSRoot());
		ArrayList<String> subErrStrs = new ArrayList<String>();
		for(Node n: nbIntNodes) {
			String subErrStr = String.format("Error: the internal node %s is non-branching.", n.getLabel());
			subErrStrs.add(subErrStr);
		}
		if(subErrStrs.size() > 0) {
			String errStr = String.format("Non-branching internal nodes found.");
			throw new IrregularTreeException(errStr, subErrStrs);
		}
	}
	public static Tree copyTree(Tree tree) {
		Tree treeCopy = new Tree(tree.getLayer());
		// IF there is a PS root, copy PS tree 
		if (tree.hasPSRoot()) {
			treeCopy.setPSRoot(copyNodeAndDescendents(tree.getPSRoot(), treeCopy));
		}
		// Otherwise just copy tokens
		else {
			for(Token token: tree.getTokenList()) {
				treeCopy.addToken(new Token(token));
			}
		}
	
		return treeCopy;
	}
	public static Node copyNodeAndDescendents(Node node, Tree tree) {
		Node nodeCopy = new Node(node);
		for (Node child: node.getPSChildren()) {
			if (child instanceof Token) {
				Token childCopy = new Token((Token) child);
				nodeCopy.addPSChild(childCopy);
				childCopy.setPSParent(nodeCopy);
				tree.addToken(childCopy);
			} else {
				Node childCopy = copyNodeAndDescendents(child, tree);
				nodeCopy.addPSChild(childCopy);
				childCopy.setPSParent(nodeCopy);
			}
		}
		return nodeCopy;
	}
}
